import zmq
import sys
import threading
import os, time
import cv2
import json
import numpy as np

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")))
from paddle_jetson import Detector


class det_task:
    def __init__(self):
        self.stop = False
        self.server = self.get_server(5002)
        self.flag_infer_initok = False
        model_path = 'paddle_jetson/det_task'
        self.yolo_predictor = Detector(model_dir=os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", model_path),
                                       device="CPU",
                                       run_mode="paddle")
        thread_tmp = threading.Thread(target=self.process_demo)
        thread_tmp.daemon = True
        thread_tmp.start()

    def init(self):
        img = np.zeros((416, 416, 3), np.uint8)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        infer_tmp = self.det
        infer_tmp(img)
        self.flag_infer_initok = True
        while True:
            try:
                if self.stop:
                    break
                time.sleep(1)
            except Exception as e:
                print(e)
                break

    def det(self, image):
        image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        det_res = self.yolo_predictor.predict_image([image_rgb], visual=False)
        det_res = self.yolo_predictor.filter_box(det_res, 0.5)
        print(det_res["boxes"])
        # 返回值处理
        return det_res["boxes"]

    def get_server(self, port):
        context = zmq.Context()
        socket = context.socket(zmq.REP)
        socket.bind(f"tcp://127.0.0.1:{port}")
        return socket

    def process_demo(self):
        print(time.strftime("%Y-%m-%d %H:%M:%S"), "{} process start".format(self.server))
        while True:
            response = self.server.recv()
            head = response[:5]
            res = []
            if head == b"Astop":
                self.stop = True
                break
            if head == b"ATATA":
                if self.flag_infer_initok:
                    res = True
                else:
                    res = False
            elif head == b"image":
                # 把bytes转为jpg格式
                img = cv2.imdecode(np.frombuffer(response[5:], dtype=np.uint8), 1)
                if self.flag_infer_initok:
                    res = self.det(img).tolist()
            json_data = json.dumps(res)
            json_data = bytes(json_data, encoding='utf-8')
            self.server.send(json_data)


main = det_task()
main.init()
